#include <position.h>
#include <movegen.h>

/*
 * Hashing scheme (starting from an empty board, white to move)
 * - xor of zobrist (piece/square)
 * - xor the side to move
 * - xor castling capabilities (bits 0-3)
 * - xor en passant square (bits 4-10)
 *
 * Pawn hash is just the zobrist xor of pawns
 *
 * Material hash is designed to be invariant with respect to piece
 * positions.  It is given by the xor of zobrist[side][piece][c],
 * where c is the number of pieces of that type.
 *
 * Material signature just contains the number of each piece type
 * (excluded king) in 5 nibbles per side.
 * - bits 0-19 white counters (0-3 pawns ...)
 * - bits 20-39 black counters
 */

#define MATERIAL_SIG_GET(m, side, type)	\
  (((m) >> (((side) * 20) + ((type) << 2))) & 15)
#define MATERIAL_SIG_INC(m, side, type)	\
  ((m) += (UINT64_C(1) << (((side) * 20) + ((type) << 2))))
#define MATERIAL_SIG_DEC(m, side, type)	\
  ((m) -= (UINT64_C(1) << (((side) * 20) + ((type) << 2))))


void position_clear(position_t* pos)
{
  assert(pos != 0);
  pos->bb_colors[WHITE] = 0;
  pos->bb_colors[BLACK] = 0;
  for (int i = 0; i < 6; i++)
    pos->bb_pieces[i] = 0;
  for (int i = 0; i < 64; i++)
    pos->squares[i] = EMPTY;
  pos->stm = WHITE;
  pos->castling_caps = (WCASTLE_SHORT | WCASTLE_LONG | BCASTLE_SHORT | BCASTLE_LONG);
  pos->enpassant = NO_ENPASSANT;
  pos->halfmoves = 0;
  pos->fullmoves = 1;
  pos->key = (pos->enpassant << 4) + pos->castling_caps;
  pos->tag = (pos->enpassant << 4) + pos->castling_caps;
  pos->pawn_key = 0;
  pos->material_key = 0;
  pos->material_signature = 0;
}


void position_drop_piece(position_t* pos, int square, int side, int piece)
{
  assert(pos != 0);
  assert(square >= 0 && square < 64);
  assert(piece >= 0 && piece < EMPTY);
  assert(side == WHITE || side == BLACK);
  assert(pos->squares[square] == EMPTY);

  BB_SET(pos->bb_colors[side], square);
  BB_SET(pos->bb_pieces[piece], square);
  pos->squares[square] = piece;
  pos->key ^= zobrist_keys[side][piece][square][0];
  pos->tag ^= zobrist_keys[side][piece][square][1];
  if (piece == PAWN)
    pos->pawn_key ^= zobrist_keys[side][piece][square][0];

  if (piece != KING) {
    int c = MATERIAL_SIG_GET(pos->material_signature, side, piece);
    pos->material_key ^= zobrist_keys[side][piece][c][0];
    MATERIAL_SIG_INC(pos->material_signature, side, piece);
    assert(MATERIAL_SIG_GET(pos->material_signature, side, piece) ==
	   BB_POPCOUNT(BB_POS_PIECES(pos, side, piece)));
  }
}


void position_remove_piece(position_t* pos, int square)
{
  assert(pos != 0);
  assert(square >= 0 && square < 64);
  assert(pos->squares[square] != EMPTY);

  int p = pos->squares[square];
  int side = position_get_square_side(pos, square);
  BB_CLEAR(pos->bb_colors[side], square);
  BB_CLEAR(pos->bb_pieces[p], square);
  pos->squares[square] = EMPTY;
  pos->key ^= zobrist_keys[side][p][square][0];
  pos->tag ^= zobrist_keys[side][p][square][1];
  if (p == PAWN)
    pos->pawn_key ^= zobrist_keys[side][p][square][0];

  if (p != KING) {
    MATERIAL_SIG_DEC(pos->material_signature, side, p);
    int c = MATERIAL_SIG_GET(pos->material_signature, side, p);
    pos->material_key ^= zobrist_keys[side][p][c][0];
    assert(MATERIAL_SIG_GET(pos->material_signature, side, p) ==
	   BB_POPCOUNT(BB_POS_PIECES(pos, side, p)));
  }
}


void position_move_piece(position_t* pos, int from, int to)
{
  assert(pos != 0);
  assert(from >= 0 && from < 64);
  assert(to >= 0 && to < 64);
  assert(pos->squares[from] != EMPTY);
  assert(pos->squares[to] == EMPTY);
  assert(from != to);

  int p = pos->squares[from];
  bitboard_t change = BB_SINGLE(from) | BB_SINGLE(to);
  int color = position_get_square_side(pos, from);
  pos->bb_colors[color] ^= change;
  pos->bb_pieces[p] ^= change;
  pos->squares[to] = p;
  pos->squares[from] = EMPTY;

  pos->key ^= zobrist_keys[color][p][from][0];
  pos->key ^= zobrist_keys[color][p][to][0];
  pos->tag ^= zobrist_keys[color][p][from][1];
  pos->tag ^= zobrist_keys[color][p][to][1];
  if (p == PAWN) {
    pos->pawn_key ^= zobrist_keys[color][p][from][0];
    pos->pawn_key ^= zobrist_keys[color][p][to][0];
  }
}


int position_get_square(const position_t* pos, int square)
{
  assert(pos != 0);
  assert(square >= 0 && square < 64);

  return pos->squares[square];
}


int position_get_square_side(const position_t* pos, int square)
{
  assert(pos != 0);
  assert(square >= 0 && square < 64);

  return ((pos->bb_colors[BLACK] >> square) & 1);
}


int position_piece_count(const position_t* pos, int side, int piece)
{
  assert(pos != 0);
  assert(piece != EMPTY);

  if (piece == KING) {
    assert(BB_POPCOUNT(BB_POS_PIECES(pos, side, KING)) == 1);
    return 1;
  }
  return MATERIAL_SIG_GET(pos->material_signature, side, piece);
}


void position_change_stm(position_t* pos)
{
  assert(pos != 0);
  pos->stm ^= 1;
  pos->key ^= side_hash_keys[0];
  pos->tag ^= side_hash_keys[1];
}


int position_get_stm(const position_t* pos)
{
  assert(pos != 0);
  return pos->stm;
}


unsigned int position_get_castling_caps(const position_t* pos)
{
  assert(pos != 0);
  return pos->castling_caps;
}


void position_set_castling_caps(position_t* pos, unsigned int caps)
{
  assert(pos != 0);
  assert(!(caps & ~(WCASTLE_SHORT | WCASTLE_LONG | BCASTLE_SHORT | BCASTLE_LONG)));

  pos->key ^= (pos->castling_caps ^ caps);
  pos->tag ^= (pos->castling_caps ^ caps);

  pos->castling_caps = caps;
}


int position_get_ep_square(const position_t* pos)
{
  assert(pos != 0);
  return pos->enpassant;
}


void position_set_ep_square(position_t* pos, int ep)
{
  assert(pos != 0);
  assert((ep >= 16 && ep < 24) || (ep >= 40 && ep < 48) || ep == NO_ENPASSANT);

  pos->key ^= ((pos->enpassant ^ ep) << 4);
  pos->tag ^= ((pos->enpassant ^ ep) << 4);

  pos->enpassant = ep;
}


int position_get_half_moves(const position_t* pos)
{
  assert(pos != 0);
  return pos->halfmoves;
}


void position_set_half_moves(position_t* pos, int n)
{
  assert(pos != 0);
  assert(n >= 0);
  pos->halfmoves = n;
}


int position_get_full_moves(const position_t* pos)
{
  assert(pos != 0);
  return pos->fullmoves;
}


void position_set_full_moves(position_t* pos, int n)
{
  assert(pos != 0);
  assert(n > 0);
  pos->fullmoves = n;
}


hash_key_t position_get_hash_key(const position_t* pos)
{
  assert(pos != 0);
  return pos->key;
}


hash_key_t position_get_hash_tag(const position_t* pos)
{
  assert(pos != 0);
  return pos->tag;
}


hash_key_t position_get_pawn_hash_key(const position_t* pos)
{
  assert(pos != 0);
  return pos->pawn_key;
}


hash_key_t position_get_material_hash_key(const position_t* pos)
{
  assert(pos != 0);
  return pos->material_key;
}


hash_key_t position_get_material_signature(const position_t* pos)
{
  assert(pos != 0);
  return pos->material_signature;
}


int position_verify_legal(const position_t* pos)
{
  /* 1) Verify that pieces do not exceed the allowed number for their
     type (taking into account possible promotions. */
  static const int maxp[6] = {8, 10, 10, 10, 9, 1};
  for (int s = 0; s < 2; s++)
    for (int p = 0; p < 6; p++)
      if (BB_POPCOUNT(BB_POS_PIECES(pos, s, p)) > maxp[p])
	return 0;

  /* 2) Verify the presence of both kings. */
  if (!BB_POS_PIECES(pos, WHITE, KING) ||
      !BB_POS_PIECES(pos, BLACK, KING))
    return 0;

  /* 3) Verify castling capabilities. */
  int caps = position_get_castling_caps(pos);
  if (caps & WCASTLE_SHORT) {
    if (!BB_TEST(BB_POS_PIECES(pos, WHITE, KING), E1) ||
	!BB_TEST(BB_POS_PIECES(pos, WHITE, ROOK), H1))
      return 0;
  }
  if (caps & WCASTLE_LONG) {
    if (!BB_TEST(BB_POS_PIECES(pos, WHITE, KING), E1) ||
	!BB_TEST(BB_POS_PIECES(pos, WHITE, ROOK), A1))
      return 0;
  }
  if (caps & BCASTLE_SHORT) {
    if (!BB_TEST(BB_POS_PIECES(pos, BLACK, KING), E8) ||
	!BB_TEST(BB_POS_PIECES(pos, BLACK, ROOK), H8))
      return 0;
  }
  if (caps & BCASTLE_LONG) {
    if (!BB_TEST(BB_POS_PIECES(pos, BLACK, KING), E8) ||
	!BB_TEST(BB_POS_PIECES(pos, BLACK, ROOK), A8))
      return 0;
  }

  /* 4) Verify the position of pawns. */
  if (BB_POS_PIECES(pos, WHITE, PAWN) & (BB_RANK_1 | BB_RANK_8))
    return 0;
  if (BB_POS_PIECES(pos, BLACK, PAWN) & (BB_RANK_1 | BB_RANK_8))
    return 0;

  /* 5) Verify ep square. */
  int stm = position_get_stm(pos);
  int ep = position_get_ep_square(pos);
  if (ep != NO_ENPASSANT) {
    if (stm == WHITE) {
      if (!BB_TEST(BB_RANK_6, ep))
	return 0;
      if (!BB_TEST(BB_POS_PIECES(pos, BLACK, PAWN), ep - 8))
	return 0;
    } else {
      if (!BB_TEST(BB_RANK_3, ep))
	return 0;
      if (!BB_TEST(BB_POS_PIECES(pos, WHITE, PAWN), ep + 8))
	return 0;
    }
  }

  /* 6) Verify that the side to move cannot capture the opponent
     king. */
  position_t temp = *pos;
  position_change_stm(&temp);
  if (movegen_detect_check(&temp))
    return 0;

  return 1;
}


void position_reverse(const position_t* src, position_t* dst)
{
  assert(src != 0);
  assert(dst != 0);
  position_t res;

  position_clear(&res);
  for (int sq = 0; sq < 64; sq++) {
    int rev = SQUARE(FILE(sq), 7 - RANK(sq));
    int p = position_get_square(src, sq);
    if (p != EMPTY) {
      int side = position_get_square_side(src, sq);
      position_drop_piece(&res, rev, 1 - side, p);
    }
  }

  if (position_get_stm(src) == position_get_stm(&res))
    position_change_stm(&res);

  unsigned int scaps = position_get_castling_caps(src);
  unsigned int dcaps = 0;
  if (scaps & WCASTLE_SHORT)
    dcaps |= BCASTLE_SHORT;
  if (scaps & WCASTLE_LONG)
    dcaps |= BCASTLE_LONG;
  if (scaps & BCASTLE_SHORT)
    dcaps |= WCASTLE_SHORT;
  if (scaps & BCASTLE_LONG)
    dcaps |= WCASTLE_LONG;
  position_set_castling_caps(&res, dcaps);

  int ep = position_get_ep_square(src);
  if (ep != NO_ENPASSANT)
    position_set_ep_square(&res, SQUARE(FILE(ep), 7 - RANK(ep)));

  position_set_half_moves(&res, position_get_half_moves(src));
  position_set_full_moves(&res, position_get_full_moves(src));

  *dst = res;
}


int position_check_hash(const position_t* pos)
{
  hash_key_t key = 0;
  hash_key_t tag = 0;

  for (int sq = 0; sq < 64; sq++) {
    int p = pos->squares[sq];
    if (p != EMPTY) {
      int s = position_get_square_side(pos, sq);
      key ^= zobrist_keys[s][p][sq][0];
      tag ^= zobrist_keys[s][p][sq][1];
    }
  }

  key ^= pos->castling_caps;
  tag ^= pos->castling_caps;

  key ^= (pos->enpassant << 4);
  tag ^= (pos->enpassant << 4);

  if (pos->stm == BLACK) {
    key ^= side_hash_keys[0];
    tag ^= side_hash_keys[1];
  }

  return (key == pos->key && tag == pos->tag);
}
